from sqlalchemy import create_engine, text
from sqlalchemy.exc import SQLAlchemyError
from config import mysql_uri
import threading

# 创建一个db的类,用于封装mysql相关的方法


class MysqlTool:
    def __init__(self) -> None:
        self.engine = create_engine(mysql_uri)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # 如何关闭mysql的链接?
        self.engine.dispose()

    def execute_sql(self, sql: str) -> bool:
        """执行传入的SQL语句，并返回是否执行成功"""
        try:
            with self.engine.connect() as connection:
                with connection.begin():  # 启用事务
                    connection.execute(text(sql))
            return True
        except SQLAlchemyError as e:
            print(f"SQL执行失败: {e}")
            return False

    def upsert_dataframe(
        self, df, table_name, engine, temp_table=None, primary_key="id"
    ):
        # 为每个线程生成唯一的临时表名，避免并发冲突
        if temp_table is None:
            temp_table = f"temp_table_{threading.get_ident()}"
        
        df.to_sql(name=temp_table, con=engine, if_exists="replace", index=False)

        # 为列名添加反引号，避免使用SQL保留字作为列名导致的问题
        columns = df.columns.tolist()
        quoted_columns = [f"`{col}`" for col in columns]  # 添加反引号
        update_statement = ", ".join(
            [f"`{col}`=VALUES(`{col}`)" for col in columns if col != primary_key]
        )

        # 使用反引号包围列名
        sql = f"""
        INSERT INTO `{table_name}` ({', '.join(quoted_columns)})
        SELECT {', '.join([f"{temp_table}.`{col}`" for col in columns])}
        FROM `{temp_table}`
        ON DUPLICATE KEY UPDATE {update_statement};
        """

        with engine.connect() as conn:
            conn.execute(text(sql))  # 使用text确保SQL语句被正确解析
            conn.execute(text(f"DROP TABLE IF EXISTS `{temp_table}`;"))

    def upsert_df(self, df, table_name, primary_key="id"):
        # table_name 必须是存在的,才能使用这个缓存表
        # logger.debug(len(df))
        self.upsert_dataframe(
            df,
            table_name,
            self.engine,
            primary_key=primary_key,
        )

    def get_max_date(self, table_name, date_column):
        """
        查询指定表中指定日期字段的最大值
        
        Args:
            table_name (str): 表名
            date_column (str): 日期字段名
        
        Returns:
            datetime or None: 最大日期值，如果表为空则返回None
        """
        try:
            sql = f"SELECT MAX({date_column}) as max_date FROM {table_name}"
            with self.engine.connect() as connection:
                result = connection.execute(text(sql))
                row = result.fetchone()
                if row and row[0]:
                    # 如果是字符串格式的日期，需要转换
                    from datetime import datetime
                    if isinstance(row[0], str):
                        return datetime.strptime(row[0][:10], '%Y-%m-%d')
                    return row[0]
                return None
        except SQLAlchemyError as e:
            print(f"查询最大日期失败: {e}")
            return None

    def delete_data_by_date_range(self, table_name, date_column, start_date, end_date):
        """
        删除指定日期范围的数据
        
        Args:
            table_name (str): 表名
            date_column (str): 日期字段名
            start_date (str): 开始日期 (YYYY-MM-DD)
            end_date (str): 结束日期 (YYYY-MM-DD)
        
        Returns:
            bool: 删除是否成功
        """
        try:
            sql = f"""
            DELETE FROM {table_name} 
            WHERE DATE({date_column}) >= '{start_date}' 
            AND DATE({date_column}) <= '{end_date}'
            """
            with self.engine.connect() as connection:
                with connection.begin():
                    result = connection.execute(text(sql))
                    deleted_count = result.rowcount
                    print(f"删除了 {deleted_count} 条记录，日期范围: {start_date} 到 {end_date}")
            return True
        except SQLAlchemyError as e:
            print(f"删除数据失败: {e}")
            return False
